import numpy as np
import os

# 创建输出目录
os.makedirs("sim_data/meta", exist_ok=True)
os.makedirs("sim_data/imu", exist_ok=True)
os.makedirs("sim_data/odom", exist_ok=True)
os.makedirs("sim_data/gps", exist_ok=True)

# ================== 1. 生成真实位姿轨迹 ==================
dt = 0.1                   # 时间步长 (s)
total_time = 50.0          # 总时长 (s)
steps = int(total_time / dt)

# 初始化轨迹（匀速圆周运动模型）
pose_gt = np.zeros((steps, 3))  # [x, y, theta]
linear_vel = 0.5                # 线速度 (m/s)
angular_vel = 0.2               # 角速度 (rad/s)

for t in range(1, steps):
    pose_gt[t, 2] = pose_gt[t-1, 2] + angular_vel * dt  # 角度更新
    pose_gt[t, 0] = pose_gt[t-1, 0] + linear_vel * np.cos(pose_gt[t, 2]) * dt  # x坐标
    pose_gt[t, 1] = pose_gt[t-1, 1] + linear_vel * np.sin(pose_gt[t, 2]) * dt  # y坐标

# 保存真值轨迹
np.savetxt("sim_data/meta/trajectory_gt.txt", pose_gt, 
           header="x[m] y[m] theta[rad]", fmt="%.6f")

# ================== 2. 生成三轴IMU数据 ==================
imu_data = []
gyro_bias = 0.005          # 陀螺仪零偏 (rad/s)
acc_noise_std = 0.1        # 加速度计噪声标准差 (m/s²)

for t in range(steps):
    # 加速度计数据（含向心加速度）
    centripetal_acc = linear_vel * angular_vel
    acc_x = centripetal_acc * np.sin(pose_gt[t, 2]) + np.random.normal(0, acc_noise_std)
    acc_y = -centripetal_acc * np.cos(pose_gt[t, 2]) + np.random.normal(0, acc_noise_std)
    #acc_z = 9.81 + np.random.normal(0, acc_noise_std)  # 重力加速度
    acc_z = 9.81  # 重力加速度
    # 陀螺仪数据（含零偏）
    gyro_z = angular_vel + gyro_bias + np.random.normal(0, 0.01)
    
    imu_data.append([t*dt, acc_x, acc_y, acc_z, 0, 0, gyro_z])

# 保存IMU数据（CSV格式）
np.savetxt("sim_data/imu/imu_data.csv", imu_data,
           header="timestamp,acc_x,acc_y,acc_z,gyro_x,gyro_y,gyro_z",
           delimiter=",", fmt="%.6f")

# ================== 3. 轮式里程计生成与积分 ==================
# 3.1 生成带噪声的原始速度测量值
odom_raw = np.zeros((steps, 3))      # [vx, vy, omega]

for t in range(steps):
    true_vx = linear_vel * np.cos(pose_gt[t, 2])
    true_vy = linear_vel * np.sin(pose_gt[t, 2])
    true_omega = angular_vel
    
    # 逐元素添加标量噪声
    odom_raw[t, 0] = true_vx + np.random.normal(0, 0.02)
    odom_raw[t, 1] = true_vy + np.random.normal(0, 0.02)
    odom_raw[t, 2] = true_omega + np.random.normal(0, 0.01)

# 保存原始里程计数据
np.savetxt("sim_data/odom/wheel_odom_raw.csv", odom_raw,
           header="vx,vy,omega", delimiter=",", fmt="%.6f")

# 3.2 轮式里程计积分（修正后的运动学模型）
odom_integrated = np.zeros((steps, 3))  # 积分后的位姿估计

for t in range(1, steps):
    # 读取当前时刻的测量值
    vx = odom_raw[t, 0]
    vy = odom_raw[t, 1]
    omega = odom_raw[t, 2]
    
    # 计算位姿增量
    delta_theta = omega * dt
    delta_x = (vx * np.cos(odom_integrated[t-1, 2]) - vy * np.sin(odom_integrated[t-1, 2])) * dt
    delta_y = (vx * np.sin(odom_integrated[t-1, 2]) + vy * np.cos(odom_integrated[t-1, 2])) * dt
    
    # 更新位姿
    odom_integrated[t, 0] = odom_integrated[t-1, 0] + delta_x
    odom_integrated[t, 1] = odom_integrated[t-1, 1] + delta_y
    odom_integrated[t, 2] = odom_integrated[t-1, 2] + delta_theta

# 保存积分轨迹
np.savetxt("sim_data/odom/wheel_odom_integrated.txt", odom_integrated,
           header="x_est[m] y_est[m] theta_est[rad]", fmt="%.6f")


# ================== 4. GPS数据生成 ==================
gps_data = []
gps_freq = 1.0                # GPS更新频率 (Hz)
gps_std = 0.3                 # GPS噪声标准差 (m)
gps_steps = int(total_time * gps_freq)

for t_gps in np.linspace(0, total_time, gps_steps, endpoint=False):
    idx = int(t_gps / dt)     # 找到对应时间戳的位姿
    
    # 添加高斯噪声
    x_gps = pose_gt[idx, 0] + np.random.normal(0, gps_std)
    y_gps = pose_gt[idx, 1] + np.random.normal(0, gps_std)
    
    gps_data.append([t_gps, x_gps, y_gps])

# 保存GPS数据（假设为局部坐标系）
np.savetxt("sim_data/gps/gps_data.csv", gps_data,
           header="timestamp,x_gps[m],y_gps[m]", delimiter=",", fmt="%.6f")

# ================== 5. 生成时间戳文件 ==================
timestamps = [t*dt for t in range(steps)]
np.savetxt("sim_data/meta/timestamps.txt", timestamps, fmt="%.6f")

